# -*- coding: utf-8 -*-
# @Function: LLM查询测试器
# @Description: 使用大模型回答测试问题，结合数据库查询和LLM生成能力，提供更智能的问答功能
# @Usage: 可独立运行或被其他模块导入，提供LLM增强的问答功能
# @Dependencies: db_manager.py, tokenizer.py, 外部LLM服务

import os
import json
from typing import List, Dict
from dotenv import load_dotenv
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_openai import ChatOpenAI

from Z_db_manager import DBManager
from Z_tokenizer import Tokenizer

# 加载环境变量
load_dotenv()

# 阿里云百炼配置
DASHSCOPE_CONFIG = {
    "api_key": os.getenv('DASHSCOPE_API_KEY'),
    "base_url": os.getenv('BASE_URL'),
    "model": "qwen-long"
}

class LLMQueryTester:
    def __init__(self):
        """初始化LLM查询测试器"""
        self.db_manager = DBManager()
        self.tokenizer = Tokenizer()
        
        # 初始化阿里云百炼模型
        self.llm = ChatOpenAI(
            api_key=DASHSCOPE_CONFIG["api_key"],
            base_url=DASHSCOPE_CONFIG["base_url"],
            model=DASHSCOPE_CONFIG["model"],
            temperature=0.7
        )
        
        # 创建回答生成的提示模板
        self.answer_template = PromptTemplate(
            input_variables=["question", "context"],
            template="""你是一个专业的水库调度规程问答助手。请根据以下参考信息，回答用户的问题。

用户问题：
{question}

参考信息：
{context}

要求：
1. 基于参考信息提供准确、完整的回答
2. 如果参考信息不足以回答问题，请明确指出
3. 回答要简洁明了，直接针对问题给出答案
4. 不要在回答中提及"根据参考信息"等字样
5. 使用专业、客观的语气

请提供你的回答：
"""
        )
        
        # 创建回答生成LLM链
        self.answer_chain = LLMChain(
            llm=self.llm,
            prompt=self.answer_template
        )

    def test_query(self, test_question: str, limit: int = 20):
        """测试查询功能并使用大模型生成回答
    
        Args:
            test_question (str): 测试问题
            limit (int, optional): 返回结果数量限制. Defaults to 20.
    
        Returns:
            dict: 查询结果和大模型回答
        """
        try:
            # 连接数据库
            self.db_manager.connect()
            
            # 对测试问题进行分词
            token_data = self.tokenizer.tokenize_question(test_question)
            tokens = token_data.get('tokens', [])
            category_id = token_data.get('category_id')  # 获取提取的category_id
            
            if not tokens:
                print("测试问题分词结果为空")
                return {"answer": "无法处理您的问题，请尝试重新表述。"}
                
            # 查询相关问题，传入category_id参数
            results = self.db_manager.query_qa(tokens, limit, category_id)
            
            # 输出查询结果
            print(f"\n测试问题: {test_question}")
            print(f"分词结果: {tokens}")
            if category_id:
                print(f"识别的水库ID: {category_id}")
            print(f"查询到 {len(results)} 个相关问题")
            
            # 准备上下文信息
            context = ""
            qa_results = []
            for i, (question, question_tokens, answer, cat_id, match_count) in enumerate(results, 1):
                answer_text = answer['answer'] if isinstance(answer, dict) else answer
                print(f"结果 {i}: {question} (类别: {cat_id}, 匹配度: {match_count})")
                print(f"答案 {i}: {answer_text}\n")
                context += f"问题 {i}: {question}\n答案 {i}: {answer_text}\n\n"
                qa_results.append({
                    "question": question,
                    "answer": answer_text,
                    "category_id": cat_id,
                    "match_count": match_count
                })
            
            # 如果没有查询结果，返回提示信息
            if not results:
                return {"answer": "抱歉，没有找到与您问题相关的信息。"}
            
            # 使用大模型生成回答
            print("\n使用大模型生成回答...")
            llm_answer = self.answer_chain.run(question=test_question, context=context)
            print(f"大模型回答: {llm_answer}")
            
            return {
                "question": test_question,
                "category_id": category_id,
                "results": qa_results,
                "answer": llm_answer
            }
            
        except Exception as e:
            print(f"查询失败: {e}")
            return {"answer": f"处理问题时出错: {str(e)}"}
        finally:
            # 关闭数据库连接
            self.db_manager.close()

def main():
    """主函数"""
    print("=== LLM问答测试工具 ===")
    
    # 验证API密钥
    if not DASHSCOPE_CONFIG["api_key"]:
        print("错误: 请设置DASHSCOPE_API_KEY环境变量")
        return
    
    print(f"连接到阿里云百炼服务: {DASHSCOPE_CONFIG['base_url']}")
    print(f"使用模型: {DASHSCOPE_CONFIG['model']}")
    
    tester = LLMQueryTester()
    
    # 测试问题列表
    test_questions = [
        "大伙房水库主坝上游边坡正常运用的设计洪水位的规定安全系数是多少？"
    ]
    
    # 交互式问答
    while True:
        print("\n" + "=" * 50)
        
        choice = "1"
        
        if choice == "0":
            break
        elif choice == "1":
            for i, question in enumerate(test_questions, 1):
                print(f"\n测试问题 {i}: {question}")
                print("-" * 50)
                result = tester.test_query(question)
                # 不再重复输出大模型回答
                print("=" * 50)
            break
        elif choice == "2":
            question = input("\n请输入您的问题: ")
            if question.strip():
                print("-" * 50)
                result = tester.test_query(question)
                # 不再重复输出大模型回答
                print("=" * 50)
        else:
            print("无效的选择，请重新输入")

if __name__ == "__main__":
    main()